import os
import sys
import importlib
def check_module(module_name):
    try:
        importlib.import_module(module_name)
    except:
        print("\033[31m[> 错误 <]\033[0m" + "acl模块未使能")
        exit(1)
check_module('acl')
import acl
import csv
import copy
import shutil
import glob
import subprocess
import numpy as np
from tbe.common.platform import set_current_compile_soc_info
from .atc_om import atc
from .constant import acl_dtype
from .profiling import op_profiling
from .dataloader import save_data, load_data
from .acl_process import acl_run

def append_to_table(file_path, data):
    import pandas as pd
    from datetime import datetime
    try:
        df = pd.read_csv(file_path, encoding='utf-8_sig')
    except FileNotFoundError:
        df = pd.DataFrame(columns=['序号', '时间'])
    current_time = datetime.now().strftime('%m-%d %H:%M')
    new_data = {
        '序号': len(df) + 1,
        '时间': current_time,
        **data
    }
    df = pd.concat([df, pd.DataFrame([new_data])], ignore_index=True)
    df.to_csv(file_path, index=False, encoding='utf-8_sig')

class Tiktools():
    def __init__(self):
        self.current_path = os.path.abspath(os.path.dirname(__file__))[:-6]
        self.prof_res = [{'task_time': None, 'aicore_time': None}]
        self.csv_data = {}
        self.is_dynamic = False
        self.save_profiling = True
        self.profiling = True
        self.one_line_num = 6
        self.no_list = ["Model ID", "Stream ID", "Task Start Time", "Input Shapes", "Input Data Types", "Input Formats", "Output Shapes", \
                        "Output Data Types", "Output Formats", "total_cycles", "mac_time(us)", "scalar_time(us)", "mte1_time(us)", \
                        "mte2_time(us)", "mte3_time(us)", "icache_miss_rate", "memory_bound"]
        self.yes_list = []
        self.op_debug_level = 1
        self.op_precision_mode = None
        self.compile_flag = 0
        self.atc_log = "debug"
        self.use_core_num = None
        self.use_ais_bench = False
        self.no_input = False
        self.deterministic = 0
        self.inplace_input_index = {}
        self.op_debug_config = ""
        self.acl_run_times = 1
        self.run_profiling = 1
        self.run_profiling_mode = "acl"  # "acl" "sim" "on"

    def init_dir(self, dir_name):
        dir_path = os.path.join(self.current_path, dir_name)
        if not os.path.exists(dir_path):
            os.mkdir(dir_path)
        return dir_path

    def clear_data(self, op_dir):
        if os.path.exists(op_dir):
            op_dir_bak = op_dir + "_bak"
            if os.path.exists(op_dir_bak):
                shutil.rmtree(op_dir_bak)
            subprocess.run("mv %s %s"%(op_dir, op_dir_bak), shell=True)

    def get_attr_type(self, value):
        dtype = None
        if isinstance(value, int):
            dtype = "int"
        if isinstance(value, float):
            dtype = "float"
        if isinstance(value, str):
            dtype = "string"
        if isinstance(value, bool):
            dtype = "bool"
        if isinstance(value, (list, tuple)):
            v = value[0]
            if isinstance(v, int):
                dtype = "list_int"
            if isinstance(v, float):
                dtype = "list_float"
            if isinstance(v, bool):
                dtype = "list_bool"
            if isinstance(v, str):
                dtype = "list_string"
            if isinstance(v, list) and isinstance(v[0], int):
                dtype = "list_list_int"
        return dtype

    def check_info(self, info):
        if "shape" not in info:
            print("\033[31m[> 错误 <]\033[0m" + "构造的输入或输出没有包含'shape', 传入的是%s" % info)
            sys.exit(1)
        if "dtype" not in info:
            print("\033[31m[> 错误 <]\033[0m" + "构造的输入或输出没有包含'dtype', 传入的是%s" % info)
            sys.exit(1)
        dtype = info["dtype"]
        if dtype == "float32":
            dtype = "float"
        if dtype not in acl_dtype:
            print("\033[31m[> 错误 <]\033[0m" + "输入的dtype有问题, 传入的是%s" % info)
            sys.exit(1)

    def process_data_info(self, data_infos):
        res = []
        new_data_infos = copy.deepcopy(data_infos)
        for info in new_data_infos:
            self.check_info(info)
            info["type"] = info["dtype"]
            if info["dtype"] == "float":
                info["type"] = "float32"
            info.pop("dtype")

            if -1 in info["shape"] or -2 in info["shape"]:
                self.is_dynamic = True
            if -1 in info["shape"] or "shape_range" not in info:
                shape_range_num = info["shape"].count(-1)
                shape_range = [(1, -1)] * shape_range_num
                info["shape_range"] = shape_range
            res.append(info)
        return res

    def get_data_info(self, params):
        inputs_info = params.get("input_desc")
        outputs_info = params.get("output_desc")
        return inputs_info, outputs_info

    def get_output_path(self, dir_path, num):
        res = []
        for i in range(num):
            data_path = os.path.join(dir_path, "output_%s.npy" % (i))
            res.append(data_path)
        return res

    def start_profiling(self, op_dir, device_id):
        if self.profiling:
            try:
                op_profiling.no_list = self.no_list
                op_profiling.yes_list = self.yes_list
                self.prof_res, values = op_profiling.op_profiling(op_dir, device_id)
                self.csv_data["性能"] = str(dict(zip(self.prof_res, values)))
                if self.prof_res:
                    op_profiling.format_print(self.prof_res, values, self.one_line_num)
            except Exception as e:
                print(e)
    
    def read_csv(data_path):
        with open(data_path, "r") as f:
            data = csv.reader(f)
            list_data = list(data)
        return list_data

    def start_aisbench(self, inputs_path, op_dir, device_id):
        om_file = glob.glob(op_dir + "/*.om")
        _inputs_path = ",".join(inputs_path)
        self.clear_data(op_dir + "/summary")
        ais_bench_cmd = "python3 -m ais_bench --model %s --input %s --device %d --output %s/summary/. --profiler 1" % (om_file[0], _inputs_path, device_id, op_dir)
        subprocess.run(ais_bench_cmd, shell=True)
        data_path = op_dir + "/summary/profiler/"
        names = os.listdir(data_path)
        for name in names:
            if "PROF" in name:
                data_path = os.path.join(data_path, name)
                data_path_id = os.path.join(data_path, "device_%s" % device_id)
                new_job_path = data_path_id + "/summary/"
                break
        names = os.listdir(new_job_path)
        for name in names:
            if "op_summary" in name:
                result_name = name
                break
        read_csv_path = os.path.join(new_job_path, result_name)
        with open(read_csv_path, "r") as f:
            data = csv.reader(f)
            csv_data = list(data)
        try:
            op_profiling.no_list = self.no_list
            op_profiling.yes_list = self.yes_list
            self.prof_res, values = op_profiling.op_profiling(op_dir, device_id)
            if self.prof_res:
                op_profiling.format_print(self.prof_res, values, self.one_line_num)
        except Exception as e:
            print(e)

    def run(self, op_type, inputs_info, inputs_data, outputs_info, attr_dict=None, device_id=0):
        soc_version = acl.get_soc_name()
        if self.run_profiling_mode == "sim" and device_id != 0:
            device_id = 0
            print("\033[93m[> 警告 <]\033[0m", "使用仿真模式,需要将device_id设置成0,工具自动将你设置为0卡")
        print("\033[93m[> 检查 <]\033[0m", soc_version, "| device_%d"%device_id, "| opp_path=%s"%os.environ["ASCEND_OPP_PATH"])
        set_current_compile_soc_info(soc_version, core_type="AiCore", aicore_num=self.use_core_num)
        os.environ["DEVICE_ID"] = str(device_id)
        dir_path = os.path.join(self.current_path, op_type)
        self.clear_data(dir_path)
        op_dir = self.init_dir(op_type)
        self.csv_data["算子"] = op_type
        inputs_path = save_data(op_dir, inputs_data)
        
        if not attr_dict:
            attr_dict = {}

        params = self.atc_op(op_type, inputs_info, outputs_info, attr_dict, soc_version=soc_version)
 
        inputs_info, outputs_info = self.get_data_info(params)
        attr = params.get("attr")
        self.csv_data["inputs_info"] = str(inputs_info)
        self.csv_data["outputs_info"] = str(outputs_info)
        self.csv_data["attr"] = str(attr)
        try:
            for _ in range(self.acl_run_times):
                acl_run(op_type, inputs_path, inputs_info, outputs_info, attr, device_id, op_dir, int(self.is_dynamic), self.inplace_input_index, self.run_profiling_mode)
                res_path = self.get_output_path(op_dir, len(outputs_info))
                result = load_data(res_path)
                
                if self.run_profiling and self.run_profiling_mode == "acl":
                    if self.use_ais_bench:
                        import ais_bench
                        self.start_aisbench(inputs_path, op_dir, device_id)
                    else:
                        self.start_profiling(op_dir, device_id)
            append_to_table(self.current_path+"/tik_tools.csv", self.csv_data)
            return result
        except Exception as e:
            print(e)
            sys.exit(1)

    def atc_op(self, op_type, inputs_info, outputs_info, attr_dict, soc_version="Ascend910"):
        op_dir = self.init_dir(op_type)
        params = {}
        attr = []
        params["op"] = op_type
        if not self.no_input:
            params["input_desc"] = self.process_data_info(inputs_info)
        params["output_desc"] = self.process_data_info(outputs_info)

        for key in attr_dict.keys():
            tmp_dict = {}
            value = attr_dict[key]
            dtype = self.get_attr_type(value)
            tmp_dict["name"] = key
            tmp_dict["value"] = value
            tmp_dict["type"] = dtype
            attr.append(tmp_dict)
        params["attr"] = attr
        atc.op_debug_level = int(self.op_debug_level)
        atc.op_precision_mode = self.op_precision_mode
        atc.atc_log = self.atc_log
        os.chdir(op_dir)

        return_code = atc.generate_om(params, op_dir, self.compile_flag, soc_version, self.deterministic, self.op_debug_config)
        os.chdir(self.current_path)
        return params

tiktools = Tiktools()